/**
 * Token计数和估算相关工具函数
 */
import { ChatMessage } from '../core/types.ts';

/**
 * 估算文本中的token数量
 * 采用简单的启发式方法：
 * - 英文: 约每4个字符1个token
 * - 中文: 约每个字符1.5个token
 * 
 * @param text 需要估算token数的文本
 * @returns 估算的token数量
 */
export function estimateTokenCount(text: string): number {
    if (!text) return 0;

    // 计算英文字符数（包括符号）
    // eslint-disable-next-line no-useless-escape
    const englishCharCount = (text.match(/[a-zA-Z0-9.,?!;:'"/\\()[\]{}<>|`~@#$%^&*_+=\-]/g) || []).length;

    // 计算中文字符数
    const chineseCharCount = (text.match(/[\u4e00-\u9fa5]/g) || []).length;

    // 英文每4个字符约1个token，中文每个字符约1.5个token
    return Math.ceil(englishCharCount / 4) + Math.ceil(chineseCharCount * 1.5);
}

/**
 * 计算消息列表中的总token数
 * 
 * @param messages 消息列表
 * @returns 总token数
 */
export function countTokensInMessages(messages: ChatMessage[]): number {
    return messages.reduce((sum, msg) => sum + estimateTokenCount(msg.content), 0);
}

/**
 * 计算单个消息的token数
 * @param message 聊天消息
 * @returns 估计的token数
 */
export function countTokensInMessage(message: ChatMessage): number {
    return estimateTokenCount(message.content);
}

/**
 * 限制消息列表以不超过指定的token数
 * 从最近的消息开始保留，直到达到token限制
 * 
 * @param messages 完整的消息列表
 * @param maxTokens 最大允许的token数
 * @param systemMessageTokens 系统消息已使用的token数（会从maxTokens中扣除）
 * @returns 截断后的消息列表
 */
export function limitMessagesToTokenCount(
    messages: ChatMessage[],
    maxTokens: number = 4000,
    systemMessageTokens: number = 0
): ChatMessage[] {
    // 计算可用的token数
    let availableTokens = maxTokens - systemMessageTokens;
    const result: ChatMessage[] = [];

    // 从最近的消息开始添加，直到达到token限制
    for (let i = messages.length - 1; i >= 0; i--) {
        const message = messages[i];
        const tokenCount = estimateTokenCount(message.content);

        if (tokenCount <= availableTokens) {
            // 添加到结果数组的开头（保持原始顺序）
            result.unshift(message);
            availableTokens -= tokenCount;
        } else {
            // 如果单条消息就超过了可用token，停止添加
            break;
        }
    }

    return result;
} 